import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
# Define the PyTorch model
class ConvNet(nn.Module):
def __init__(self):
super(ConvNet, self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(3, 32, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2)
)
self.layer2 = nn.Sequential(
nn.Conv2d(32, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2)
)
self.flatten = nn.Flatten()
self.fc1 = nn.Linear(64 * (IMAGE_SIZE // 4) * (IMAGE_SIZE // 4), IMAGE_SIZE)
self.fc2 = nn.Linear(IMAGE_SIZE, 8)
def forward(self, x):
x = self.layer1(x)
x = self.layer2(x)
x = self.flatten(x)
x = self.fc1(x)
x = torch.relu(x)
x = self.fc2(x)
return x
# Create the model
model = ConvNet()
# Define the loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# print(device)
# model.to(device)
def train_model(model, dataloader, criterion, optimizer, num_epochs=10):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.train()
for epoch in range(1, num_epochs + 1):
running_loss = 0.0
correct = 0
total = 0
for batch in tqdm(dataloader, desc=f"Epoch {epoch} - Training"):
images, labels = batch['image'].to(device), batch['label'].to(device) # Adjusted this line
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
labels_indices = torch.argmax(labels, dim=1)
loss = criterion(outputs, labels_indices)
# Calculate training accuracy
_, predicted = torch.max(outputs, 1)
correct += (predicted == labels_indices).sum().item()
total += labels.size(0)
train_loss = running_loss / len(dataloader)
train_accuracy = correct / total
val_loss, val_accuracy = evaluate_model(model, val_dataloader, criterion)
epoch_loss = running_loss / len(dataloader)
print(f"Epoch {epoch}: Train Loss: {train_loss:.4f}, Train Accuracy: {train_accuracy:.4f}, Val Loss: {val_loss:.4f}, Val Accuracy: {val_accuracy:.4f}")
